package com.mybatis.plugin;

import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.SystemMetaObject;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Map;
import java.util.Properties;

/**
 * 分页拦截器
 * 目标方法: org.apache.ibatis.executor.statement.StatementHandler#prepare(java.sql.Connection, java.lang.Integer)
 */
@Intercepts(@Signature(type = StatementHandler.class, method = "prepare",
        args = {Connection.class, Integer.class}))
public class PagePluginInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 检测当前查询是否带分页参数
        StatementHandler target = (StatementHandler) invocation.getTarget();
        BoundSql boundSql = target.getBoundSql();
        Object parameterObject = boundSql.getParameterObject();
        Page page = null;
        if (parameterObject instanceof Page) {
            page = (Page) parameterObject;
        } else if (parameterObject instanceof Map) {
            page = (Page) ((Map) parameterObject).values().stream().filter(v -> v instanceof Page).findFirst().orElse(null);
        }

        if (page == null) {
            return invocation.proceed();
        }

        // 设置总行数
        page.setTotal(selectCount(invocation));

        // 修改原有的SQL
        String newSql = String.format("%s limit %s offset %s", boundSql.getSql(),
                page.getSize(), page.getOffset());
        SystemMetaObject.forObject(boundSql).setValue("sql", newSql);

        return invocation.proceed();
    }

    private int selectCount(Invocation invocation) throws SQLException {
        int count = 0;
        StatementHandler target = (StatementHandler) invocation.getTarget();
        // SQL包包含SQL、参数、参数映射
        BoundSql boundSql = target.getBoundSql();
        String countSql = String.format("select count(*) from (%s) as _page", boundSql.getSql());

        // JDBC
        Connection connection = (Connection) invocation.getArgs()[0];

        PreparedStatement preparedStatement = connection.prepareStatement(countSql);
        target.getParameterHandler().setParameters(preparedStatement);

        ResultSet resultSet = preparedStatement.executeQuery();
        if (resultSet.next()) {
            count = resultSet.getInt(1);
        }
        resultSet.close();
        preparedStatement.close();
        return count;
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    @Override
    public void setProperties(Properties properties) {

    }
}
